{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "7f0195c2-4a20-488e-8782-ca5a83488d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_hub.tools.waii import WaiiToolSpec\n",
    "\n",
    "waii_tool = WaiiToolSpec(\n",
    "    url=\"https://tweakit.waii.ai/api/\",\n",
    "    # API Key of Waii (not OpenAI API key)\n",
    "    api_key=\"3........\",\n",
    "    # Which database you want to use, you need add the db connection to Waii first\n",
    "    database_key=\"snowflake://....\",\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0a79a9fa-e5ff-4242-99a2-08cc85e158a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "."
     ]
    },
    {
     "data": {
      "text/plain": [
       "'SELECT\\n    table_schema,\\n    table_name,\\n    COUNT(column_name) AS number_of_columns\\nFROM waii.information_schema.columns\\nGROUP BY\\n    table_schema,\\n    table_name\\nORDER BY\\n    table_schema,\\n    table_name\\n'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ".."
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>TABLE_SCHEMA</th>\n",
       "      <th>TABLE_NAME</th>\n",
       "      <th>NUMBER_OF_COLUMNS</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>BATTLE_DEATH</td>\n",
       "      <td>BATTLE</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>BATTLE_DEATH</td>\n",
       "      <td>DEATH</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>BATTLE_DEATH</td>\n",
       "      <td>SHIP</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CAR</td>\n",
       "      <td>CARS_DATA</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CAR</td>\n",
       "      <td>CAR_MAKERS</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>107</th>\n",
       "      <td>VOTER</td>\n",
       "      <td>VOTES</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>108</th>\n",
       "      <td>WORLD</td>\n",
       "      <td>CITY</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>109</th>\n",
       "      <td>WORLD</td>\n",
       "      <td>COUNTRY</td>\n",
       "      <td>15</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>110</th>\n",
       "      <td>WORLD</td>\n",
       "      <td>COUNTRYLANGUAGE</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>111</th>\n",
       "      <td>WORLD</td>\n",
       "      <td>SQLITE_SEQUENCE</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>112 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     TABLE_SCHEMA       TABLE_NAME  NUMBER_OF_COLUMNS\n",
       "0    BATTLE_DEATH           BATTLE                  6\n",
       "1    BATTLE_DEATH            DEATH                  5\n",
       "2    BATTLE_DEATH             SHIP                  7\n",
       "3             CAR        CARS_DATA                  8\n",
       "4             CAR       CAR_MAKERS                  4\n",
       "..            ...              ...                ...\n",
       "107         VOTER            VOTES                  5\n",
       "108         WORLD             CITY                  5\n",
       "109         WORLD          COUNTRY                 15\n",
       "110         WORLD  COUNTRYLANGUAGE                  4\n",
       "111         WORLD  SQLITE_SEQUENCE                  2\n",
       "\n",
       "[112 rows x 3 columns]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\"The table 'COLUMNS' contains the most columns. The top 5 tables with the number of columns are 'COLUMNS' with 43 columns, 'TABLES' with 25 columns, and the remaining tables have fewer than 25 columns.\""
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from llama_index import VectorStoreIndex\n",
    "\n",
    "# Use as Data Loader, load data to index and query it\n",
    "documents = waii_tool.load_data(\"Get all tables with their number of columns\")\n",
    "index = VectorStoreIndex.from_documents(documents).as_query_engine()\n",
    "\n",
    "index.query(\n",
    "    \"Which table contains most columns, tell me top 5 tables with number of columns?\"\n",
    ").response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "b259d9cd-bbb8-4fff-a4ce-80fb0f3a1a10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use as tool, initialize it\n",
    "from llama_index.agent import OpenAIAgent\n",
    "from llama_index.llms import OpenAI\n",
    "\n",
    "agent = OpenAIAgent.from_tools(\n",
    "    waii_tool.to_tool_list(), llm=OpenAI(model=\"gpt-4-1106-preview\"), verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "094b7878-59d6-4f12-b357-4f0d254953da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STARTING TURN 1\n",
      "---------------\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: get_answer with args: {\"ask\":\"What are the top 3 countries with the highest number of car factories?\"}\n",
      "."
     ]
    },
    {
     "data": {
      "text/plain": [
       "'WITH carmakerscountries AS (\\n    SELECT\\n        countryname,\\n        COUNT(id) AS numberoffactories\\n    FROM waii.car.car_makers\\n    INNER JOIN waii.car.countries\\n        ON country = countryid\\n    GROUP BY\\n        countryname\\n)\\n\\nSELECT\\n    countryname,\\n    numberoffactories\\nFROM carmakerscountries\\nORDER BY\\n    numberoffactories DESC\\nLIMIT 3\\n'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "."
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>COUNTRYNAME</th>\n",
       "      <th>NUMBEROFFACTORIES</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>japan</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>germany</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>usa</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  COUNTRYNAME  NUMBEROFFACTORIES\n",
       "0       japan                  5\n",
       "1     germany                  4\n",
       "2         usa                  4"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got output: Japan has 5 factories, Germany has 4 factories, and the USA also has 4 factories.\n",
      "========================\n",
      "\n",
      "STARTING TURN 2\n",
      "---------------\n",
      "\n",
      "The top 3 countries with the highest number of car factories are:\n",
      "\n",
      "1. Japan with 5 factories\n",
      "2. Germany with 4 factories\n",
      "3. USA with 4 factories\n"
     ]
    }
   ],
   "source": [
    "# Ask simple questions\n",
    "print(agent.chat(\"Give me top 3 countries with the most number of car factory\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "233deb3d-547b-49a2-89a4-28fa1abea9dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STARTING TURN 1\n",
      "---------------\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: get_answer with args: {\"ask\": \"What are the car factories in Japan?\"}\n",
      "."
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"SELECT DISTINCT fullname\\nFROM waii.car.car_makers\\nINNER JOIN waii.car.countries\\n    ON country = countryid\\nWHERE\\n    countryname ILIKE '%japan%'\\n\""
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "."
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>FULLNAME</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Nissan Motors</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Honda</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Mazda</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Subaru</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Toyota</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        FULLNAME\n",
       "0  Nissan Motors\n",
       "1          Honda\n",
       "2          Mazda\n",
       "3         Subaru\n",
       "4         Toyota"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got output: The result of the information from multiple sources is a list of car manufacturers. The list includes Nissan Motors, Honda, Mazda, Subaru, and Toyota.\n",
      "========================\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: get_answer with args: {\"ask\": \"What are the car factories in Germany?\"}\n",
      "."
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"SELECT maker\\nFROM waii.car.car_makers\\nINNER JOIN waii.car.countries\\n    ON country = countryid\\nWHERE\\n    countryname ILIKE '%germany%'\\n\""
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "."
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>MAKER</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>volkswagen</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>bmw</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>daimler benz</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>opel</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          MAKER\n",
       "0    volkswagen\n",
       "1           bmw\n",
       "2  daimler benz\n",
       "3          opel"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got output: The result of the information from multiple sources is that there are four different car makers mentioned: Volkswagen, BMW, Daimler Benz, and Opel.\n",
      "========================\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: get_answer with args: {\"ask\": \"What are the car factories in the USA?\"}\n",
      "."
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"SELECT DISTINCT maker\\nFROM waii.car.car_makers\\nINNER JOIN waii.car.countries\\n    ON country = countryid\\nWHERE\\n    countryname ILIKE '%usa%'\\n\""
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "."
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>MAKER</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>amc</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>gm</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>ford</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>chrysler</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      MAKER\n",
       "0       amc\n",
       "1        gm\n",
       "2      ford\n",
       "3  chrysler"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got output: The result of the information from multiple sources is that there are four different car makers mentioned: AMC, GM, Ford, and Chrysler.\n",
      "========================\n",
      "\n",
      "STARTING TURN 2\n",
      "---------------\n",
      "\n",
      "Here are the car factories in the top 3 countries with the most factories:\n",
      "\n",
      "**Japan:**\n",
      "- Nissan Motors\n",
      "- Honda\n",
      "- Mazda\n",
      "- Subaru\n",
      "- Toyota\n",
      "\n",
      "**Germany:**\n",
      "- Volkswagen\n",
      "- BMW\n",
      "- Daimler Benz\n",
      "- Opel\n",
      "\n",
      "**USA:**\n",
      "- AMC\n",
      "- GM (General Motors)\n",
      "- Ford\n",
      "- Chrysler\n"
     ]
    }
   ],
   "source": [
    "print(agent.chat(\"What are the car factories of these countries\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "90c2ba4d-6ac4-4cbb-93b0-e03a8d015042",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STARTING TURN 1\n",
      "---------------\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: get_answer with args: {\"ask\":\"What are the top 3 longest running queries including their query_id and duration?\"}\n",
      "."
     ]
    },
    {
     "data": {
      "text/plain": [
       "'SELECT\\n    query_id,\\n    total_elapsed_time\\nFROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY(RESULT_LIMIT => 10000))\\nORDER BY\\n    total_elapsed_time DESC NULLS LAST\\nLIMIT 3\\n'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "......."
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>QUERY_ID</th>\n",
       "      <th>TOTAL_ELAPSED_TIME</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>01b0c491-0001-ec2a-0022-ba8700c79812</td>\n",
       "      <td>27117</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>01b0c5df-0001-ebd9-0022-ba8700c7d18a</td>\n",
       "      <td>15903</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>01b0cb83-0001-edf6-0022-ba8700c87192</td>\n",
       "      <td>11813</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                               QUERY_ID  TOTAL_ELAPSED_TIME\n",
       "0  01b0c491-0001-ec2a-0022-ba8700c79812               27117\n",
       "1  01b0c5df-0001-ebd9-0022-ba8700c7d18a               15903\n",
       "2  01b0cb83-0001-edf6-0022-ba8700c87192               11813"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got output: The result from multiple sources shows the total elapsed time for each query. The first query, with the ID '01b0c491-0001-ec2a-0022-ba8700c79812', had a total elapsed time of 27117. The second query, with the ID '01b0c5df-0001-ebd9-0022-ba8700c7d18a', had a total elapsed time of 15903. Finally, the third query, with the ID '01b0cb83-0001-edf6-0022-ba8700c87192', had a total elapsed time of 11813.\n",
      "========================\n",
      "\n",
      "STARTING TURN 2\n",
      "---------------\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: performance_analyze with args: {\"query_uuid\": \"01b0c491-0001-ec2a-0022-ba8700c79812\"}\n",
      "Got output: {\n",
      "  \"summary\": [\n",
      "    \"The 'tablescan' operator on 'tweakit_playground.retail_data.store_sales' is the most time-consuming, with an overall percentage of execution time of 0.78, processing time of 0.01, and remote disk IO of 0.76. It scanned 2,731,773,952 bytes and emitted 110,005,394 output rows.\",\n",
      "    \"The 'aggregate' operator associated with the 'GROUP BY ss_customer_sk, d_year, d_moy' clause and SUM function is the second most time-consuming, with an overall percentage of execution time of 0.08 and processing time of 0.08. It processed 110,005,394 input rows and emitted 90,106,053 output rows.\"\n",
      "  ],\n",
      "  \"recommendations\": [\n",
      "    \"Consider merging the 'monthly_spend' and 'first_last_spend' CTEs into a single CTE using conditional aggregation. This would reduce the number of table scans and group by operations, as well as the amount of data being processed in the 'aggregate' operator.\",\n",
      "    \"Push the filter 'd_year IN (2000, 2001)' directly into the table scan of 'store_sales' to reduce the number of rows scanned and processed. This can be done by adding a WHERE clause that filters on 'd_year' in the join condition with 'date_dim'.\",\n",
      "    \"Review the need for the self-join in the final SELECT statement. If the logic can be preserved, consider using a single pass through the data to calculate the spend increase for both years, which could eliminate the need for the self-join and reduce the complexity of the query.\"\n",
      "  ],\n",
      "  \"query_text\": \"WITH monthly_spend AS (\\n    SELECT\\n        ss_customer_sk,\\n        d_year,\\n        d_moy,\\n        SUM((\\n            ss_sales_price - ss_ext_discount_amt\\n        ) * ss_quantity) AS monthly_spend\\n    FROM tweakit_playground.retail_data.store_sales\\n    INNER JOIN tweakit_playground.retail_data.date_dim\\n        ON ss_sold_date_sk = d_date_sk\\n    WHERE\\n        d_year IN (2000, 2001)\\n    GROUP BY\\n        ss_customer_sk,\\n        d_year,\\n        d_moy\\n),\\n\\nfirst_last_spend AS (\\n    SELECT\\n        ss_customer_sk,\\n        d_year,\\n        MAX(CASE WHEN d_moy = 1 THEN monthly_spend ELSE 0 END) AS first_month_spend,\\n        MAX(CASE WHEN d_moy = 12 THEN monthly_spend ELSE 0 END) AS last_month_spend\\n    FROM monthly_spend\\n    GROUP BY\\n        ss_customer_sk,\\n        d_year\\n),\\n\\nspend_increase AS (\\n    SELECT\\n        ss_customer_sk,\\n        d_year,\\n        last_month_spend - first_month_spend AS spend_increase\\n    FROM first_last_spend\\n)\\n\\nSELECT a.ss_customer_sk\\nFROM spend_increase AS a\\nINNER JOIN spend_increase AS b\\n    ON a.ss_customer_sk = b.ss_customer_sk\\nWHERE\\n    a.d_year = 2001 AND b.d_year = 2000 AND a.spend_increase > b.spend_increase\",\n",
      "  \"execution_time_ms\": 26799,\n",
      "  \"compilation_time_ms\": 318\n",
      "}\n",
      "========================\n",
      "\n",
      "STARTING TURN 3\n",
      "---------------\n",
      "\n",
      "The top 3 longest running queries are as follows:\n",
      "\n",
      "1. Query ID: `01b0c491-0001-ec2a-0022-ba8700c79812` with a duration of 27,117 ms.\n",
      "2. Query ID: `01b0c5df-0001-ebd9-0022-ba8700c7d18a` with a duration of 15,903 ms.\n",
      "3. Query ID: `01b0cb83-0001-edf6-0022-ba8700c87192` with a duration of 11,813 ms.\n",
      "\n",
      "The performance analysis of the first query (`01b0c491-0001-ec2a-0022-ba8700c79812`) reveals that the most time-consuming operation is a table scan on `tweakit_playground.retail_data.store_sales`, which accounts for 78% of the execution time. It scanned over 2.7 billion bytes and emitted 110 million rows. The second most time-consuming operation is an aggregate associated with a `GROUP BY` clause and `SUM` function, taking up 8% of the execution time and processing 110 million input rows.\n",
      "\n",
      "The recommendations to improve the performance of this query include:\n",
      "\n",
      "1. Merging the `monthly_spend` and `first_last_spend` CTEs into a single CTE using conditional aggregation to reduce the number of table scans and group by operations.\n",
      "2. Pushing the filter `d_year IN (2000, 2001)` directly into the table scan of `store_sales` to reduce the number of rows scanned.\n",
      "3. Reviewing the need for the self-join in the final `SELECT` statement and considering a single pass through the data to calculate the spend increase for both years, which could eliminate the need for the self-join.\n",
      "\n",
      "The query text is as follows:\n",
      "\n",
      "```sql\n",
      "WITH monthly_spend AS (\n",
      "    SELECT\n",
      "        ss_customer_sk,\n",
      "        d_year,\n",
      "        d_moy,\n",
      "        SUM((\n",
      "            ss_sales_price - ss_ext_discount_amt\n",
      "        ) * ss_quantity) AS monthly_spend\n",
      "    FROM tweakit_playground.retail_data.store_sales\n",
      "    INNER JOIN tweakit_playground.retail_data.date_dim\n",
      "        ON ss_sold_date_sk = d_date_sk\n",
      "    WHERE\n",
      "        d_year IN (2000, 2001)\n",
      "    GROUP BY\n",
      "        ss_customer_sk,\n",
      "        d_year,\n",
      "        d_moy\n",
      "),\n",
      "\n",
      "first_last_spend AS (\n",
      "    SELECT\n",
      "        ss_customer_sk,\n",
      "        d_year,\n",
      "        MAX(CASE WHEN d_moy = 1 THEN monthly_spend ELSE 0 END) AS first_month_spend,\n",
      "        MAX(CASE WHEN d_moy = 12 THEN monthly_spend ELSE 0 END) AS last_month_spend\n",
      "    FROM monthly_spend\n",
      "    GROUP BY\n",
      "        ss_customer_sk,\n",
      "        d_year\n",
      "),\n",
      "\n",
      "spend_increase AS (\n",
      "    SELECT\n",
      "        ss_customer_sk,\n",
      "        d_year,\n",
      "        last_month_spend - first_month_spend AS spend_increase\n",
      "    FROM first_last_spend\n",
      ")\n",
      "\n",
      "SELECT a.ss_customer_sk\n",
      "FROM spend_increase AS a\n",
      "INNER JOIN spend_increase AS b\n",
      "    ON a.ss_customer_sk = b.ss_customer_sk\n",
      "WHERE\n",
      "    a.d_year = 2001 AND b.d_year = 2000 AND a.spend_increase > b.spend_increase\n",
      "```\n",
      "\n",
      "The execution time of this query was 26,799 ms, and the compilation time was 318 ms.\n"
     ]
    }
   ],
   "source": [
    "# Do performance analysis\n",
    "print(\n",
    "    agent.chat(\n",
    "        \"Give me top 3 longest running queries, include the complete query_id and their duration. And analyze performance of the first query\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "47530eba-24be-42d9-b1a0-fa1af28934f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STARTING TURN 1\n",
      "---------------\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: diff_query with args: {\n",
      "  \"previous_query\": \"SELECT\\n    employee_id,\\n    department,\\n    salary,\\n    AVG(salary) OVER (PARTITION BY department) AS department_avg_salary,\\n    salary - AVG(salary) OVER (PARTITION BY department) AS diff_from_avg\\nFROM\\n    employees;\",\n",
      "  \"current_query\": \"SELECT\\n    employee_id,\\n    department,\\n    salary,\\n    MAX(salary) OVER (PARTITION BY department) AS department_max_salary,\\n    salary - AVG(salary) OVER (PARTITION BY department) AS diff_from_avg\\nFROM\\n    employees;\\nLIMIT 100;\"\n",
      "}\n",
      "Got output: The query retrieves the maximum salary and the difference from the average salary for each employee's department. It does this by retrieving the employee ID, department, and salary from the employees table. Then, it calculates the maximum salary for each department using the MAX() function and the PARTITION BY clause. Next, it calculates the difference between each employee's salary and the average salary for their department using the AVG() function and the PARTITION BY clause. The results from these steps are combined into a single result set. Finally, the result set is limited to the first 100 rows.\n",
      "========================\n",
      "\n",
      "STARTING TURN 2\n",
      "---------------\n",
      "\n",
      "The difference between the two queries is in the calculation of the departmental salary metric:\n",
      "\n",
      "1. The first query calculates the average salary for each department (`department_avg_salary`) and the difference between each employee's salary and the department's average salary (`diff_from_avg`).\n",
      "\n",
      "2. The second query calculates the maximum salary for each department (`department_max_salary`) instead of the average. However, it still calculates the difference between each employee's salary and the department's average salary (`diff_from_avg`), just like the first query.\n",
      "\n",
      "Additionally, the second query has a `LIMIT 100` clause, which restricts the result set to the first 100 rows. The first query does not have this limit and would return all rows from the `employees` table.\n"
     ]
    }
   ],
   "source": [
    "# Diff two queries\n",
    "previous_query = \"\"\"\n",
    "SELECT\n",
    "    employee_id,\n",
    "    department,\n",
    "    salary,\n",
    "    AVG(salary) OVER (PARTITION BY department) AS department_avg_salary,\n",
    "    salary - AVG(salary) OVER (PARTITION BY department) AS diff_from_avg\n",
    "FROM\n",
    "    employees;\n",
    "\"\"\"\n",
    "current_query = \"\"\"\n",
    "SELECT\n",
    "    employee_id,\n",
    "    department,\n",
    "    salary,\n",
    "    MAX(salary) OVER (PARTITION BY department) AS department_max_salary,\n",
    "    salary - AVG(salary) OVER (PARTITION BY department) AS diff_from_avg\n",
    "FROM\n",
    "    employees;\n",
    "LIMIT 100;\n",
    "\"\"\"\n",
    "print(agent.chat(f\"tell me difference between {previous_query} and {current_query}\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "7a222df4-d00e-4fe8-be8e-9efdb43f1462",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STARTING TURN 1\n",
      "---------------\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: describe_dataset with args: {\"ask\":\"Summarize the dataset\"}\n",
      "Got output: The dataset consists of multiple schemas, each containing tables related to different domains. The \"FLIGHT\" schema contains tables related to airlines, airports, and flights. It provides information about different airlines, airports, and flight details. The \"STUDENT_TRANSCRIPTS_TRACKING\" schema contains tables related to student information, courses, degree programs, departments, semesters, and transcripts. It provides a comprehensive database for tracking and managing student records, enrollment, and academic information. The \"WORLD\" schema contains tables related to cities, countries, and languages. It provides information about cities, countries, their populations, languages spoken, and various other attributes. The \"INFORMATION_SCHEMA\" schema provides information about the database objects and metadata in the WAII database. It includes tables such as APPLICABLE_ROLES, CLASSES, COLUMNS, DATABASES, FUNCTIONS, LOAD_HISTORY, OBJECT_PRIVILEGES, PROCEDURES, REFERENTIAL_CONSTRAINTS, SCHEMATA, SEQUENCES, and more. The \"EMPLOYEE_HIRE_EVALUATION\" schema contains tables related to employee information, evaluations, hiring evaluations, and shop details. It provides a comprehensive view of employee data, bonus awards, hiring evaluations, and shop performance. The \"PETS\" schema contains tables related to pet ownership and student information. It includes tables for tracking the relationship between students and their pets, storing pet data, and managing student records. Overall, the dataset covers a wide range of domains including airlines, airports, flights, student records, courses, degrees, cities, countries, languages, employee information, evaluations, hiring evaluations, shops, pet ownership, and student information. It provides a comprehensive database for analyzing and generating reports on various aspects of these domains.\n",
      "========================\n",
      "\n",
      "STARTING TURN 2\n",
      "---------------\n",
      "\n",
      "The dataset is a comprehensive collection of information spanning multiple domains:\n",
      "\n",
      "- **FLIGHT Schema**: Contains details about airlines, airports, and flights, including flight schedules, aircraft, and airline operations.\n",
      "\n",
      "- **STUDENT_TRANSCRIPTS_TRACKING Schema**: Manages student records, courses, degree programs, departments, semesters, and transcripts, serving as a database for educational institutions.\n",
      "\n",
      "- **WORLD Schema**: Offers data on cities, countries, and languages, including demographic and geographic information.\n",
      "\n",
      "- **INFORMATION_SCHEMA**: Holds metadata about the database itself, such as details about tables, columns, roles, and privileges.\n",
      "\n",
      "- **EMPLOYEE_HIRE_EVALUATION Schema**: Stores data on employee evaluations, hiring processes, and shop performance, useful for human resources and business analysis.\n",
      "\n",
      "- **PETS Schema**: Relates to pet ownership and student information, tracking the relationship between students and their pets.\n",
      "\n",
      "Overall, the dataset provides a rich source of data for analysis and reporting across various sectors, including aviation, education, geography, human resources, and pet ownership.\n"
     ]
    }
   ],
   "source": [
    "# Describe dataset\n",
    "print(agent.chat(\"Summarize the dataset\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "d6bdf837-241a-4637-a07e-a73fafd52a07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STARTING TURN 1\n",
      "---------------\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: transcode with args: {\"instruction\":\"Translate this PySpark query to Snowflake SQL.\",\"source_dialect\":\"spark\",\"source_query\":\"from pyspark.sql import SparkSession\\nfrom pyspark.sql.functions import avg, lag, lead, round\\nfrom pyspark.sql.window import Window\\n\\nspark = SparkSession.builder.appName(\\\"yearly_car_analysis\\\").getOrCreate()\\n\\nyearly_avg_hp = cars_data.groupBy(\\\"year\\\").agg(avg(\\\"horsepower\\\").alias(\\\"avg_horsepower\\\"))\\n\\nwindowSpec = Window.orderBy(\\\"year\\\")\\n\\nyearly_comparisons = yearly_avg_hp.select(\\n    \\\"year\\\",\\n    \\\"avg_horsepower\\\",\\n    lag(\\\"avg_horsepower\\\").over(windowSpec).alias(\\\"prev_year_hp\\\"),\\n    lead(\\\"avg_horsepower\\\").over(windowSpec).alias(\\\"next_year_hp\\\")\\n)\\n\\nfinal_result = yearly_comparisons.select(\\n    \\\"year\\\",\\n    \\\"avg_horsepower\\\",\\n    round(\\n        (yearly_comparisons.avg_horsepower - yearly_comparisons.prev_year_hp) / \\n        yearly_comparisons.prev_year_hp * 100, 2\\n    ).alias(\\\"percentage_diff_prev_year\\\"),\\n    round(\\n        (yearly_comparisons.next_year_hp - yearly_comparisons.avg_horsepower) / \\n        yearly_comparisons.avg_horsepower * 100, 2\\n    ).alias(\\\"percentage_diff_next_year\\\")\\n).orderBy(\\\"year\\\")\\n\\nfinal_result.show()\",\"target_dialect\":\"snowflake\"}\n",
      "Got output: WITH yearly_avg_hp AS (\n",
      "    SELECT\n",
      "        year,\n",
      "        AVG(horsepower) AS avg_horsepower\n",
      "    FROM waii.car.cars_data\n",
      "    GROUP BY\n",
      "        year\n",
      "),\n",
      "\n",
      "yearly_comparisons AS (\n",
      "    SELECT\n",
      "        year,\n",
      "        avg_horsepower,\n",
      "        LAG(avg_horsepower) OVER (ORDER BY year) AS prev_year_hp,\n",
      "        LEAD(avg_horsepower) OVER (ORDER BY year) AS next_year_hp\n",
      "    FROM yearly_avg_hp\n",
      ")\n",
      "\n",
      "SELECT\n",
      "    year,\n",
      "    avg_horsepower,\n",
      "    ROUND((\n",
      "        (\n",
      "            avg_horsepower - prev_year_hp\n",
      "        ) / NULLIF(prev_year_hp, 0) * 100\n",
      "    ), 2) AS percentage_diff_prev_year,\n",
      "    ROUND((\n",
      "        (\n",
      "            next_year_hp - avg_horsepower\n",
      "        ) / NULLIF(avg_horsepower, 0) * 100\n",
      "    ), 2) AS percentage_diff_next_year\n",
      "FROM yearly_comparisons\n",
      "ORDER BY\n",
      "    year\n",
      "\n",
      "========================\n",
      "\n",
      "STARTING TURN 2\n",
      "---------------\n",
      "\n",
      "The translated Snowflake SQL query from the provided PySpark code is as follows:\n",
      "\n",
      "```sql\n",
      "WITH yearly_avg_hp AS (\n",
      "    SELECT\n",
      "        year,\n",
      "        AVG(horsepower) AS avg_horsepower\n",
      "    FROM waii.car.cars_data\n",
      "    GROUP BY\n",
      "        year\n",
      "),\n",
      "\n",
      "yearly_comparisons AS (\n",
      "    SELECT\n",
      "        year,\n",
      "        avg_horsepower,\n",
      "        LAG(avg_horsepower) OVER (ORDER BY year) AS prev_year_hp,\n",
      "        LEAD(avg_horsepower) OVER (ORDER BY year) AS next_year_hp\n",
      "    FROM yearly_avg_hp\n",
      ")\n",
      "\n",
      "SELECT\n",
      "    year,\n",
      "    avg_horsepower,\n",
      "    ROUND((\n",
      "        (\n",
      "            avg_horsepower - prev_year_hp\n",
      "        ) / NULLIF(prev_year_hp, 0) * 100\n",
      "    ), 2) AS percentage_diff_prev_year,\n",
      "    ROUND((\n",
      "        (\n",
      "            next_year_hp - avg_horsepower\n",
      "        ) / NULLIF(avg_horsepower, 0) * 100\n",
      "    ), 2) AS percentage_diff_next_year\n",
      "FROM yearly_comparisons\n",
      "ORDER BY\n",
      "    year\n",
      "```\n",
      "\n",
      "This query calculates the average horsepower (`avg_horsepower`) for each year from the `cars_data` table, and then computes the percentage difference in average horsepower compared to the previous and next years. The `NULLIF` function is used to avoid division by zero. The results are ordered by year.\n"
     ]
    }
   ],
   "source": [
    "q = \"\"\"\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.functions import avg, lag, lead, round\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "spark = SparkSession.builder.appName(\"yearly_car_analysis\").getOrCreate()\n",
    "\n",
    "yearly_avg_hp = cars_data.groupBy(\"year\").agg(avg(\"horsepower\").alias(\"avg_horsepower\"))\n",
    "\n",
    "windowSpec = Window.orderBy(\"year\")\n",
    "\n",
    "yearly_comparisons = yearly_avg_hp.select(\n",
    "    \"year\",\n",
    "    \"avg_horsepower\",\n",
    "    lag(\"avg_horsepower\").over(windowSpec).alias(\"prev_year_hp\"),\n",
    "    lead(\"avg_horsepower\").over(windowSpec).alias(\"next_year_hp\")\n",
    ")\n",
    "\n",
    "final_result = yearly_comparisons.select(\n",
    "    \"year\",\n",
    "    \"avg_horsepower\",\n",
    "    round(\n",
    "        (yearly_comparisons.avg_horsepower - yearly_comparisons.prev_year_hp) / \n",
    "        yearly_comparisons.prev_year_hp * 100, 2\n",
    "    ).alias(\"percentage_diff_prev_year\"),\n",
    "    round(\n",
    "        (yearly_comparisons.next_year_hp - yearly_comparisons.avg_horsepower) / \n",
    "        yearly_comparisons.avg_horsepower * 100, 2\n",
    "    ).alias(\"percentage_diff_next_year\")\n",
    ").orderBy(\"year\")\n",
    "\n",
    "final_result.show()\n",
    "\"\"\"\n",
    "print(agent.chat(f\"translate this pyspark query {q}, to Snowflake\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "e94a9d08-1562-47c1-8983-77270bcca21e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STARTING TURN 1\n",
      "---------------\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: run_query with args: {\"sql\":\"WITH yearly_avg_hp AS (\\n    SELECT\\n        year,\\n        AVG(horsepower) AS avg_horsepower\\n    FROM waii.car.cars_data\\n    GROUP BY\\n        year\\n),\\n\\nyearly_comparisons AS (\\n    SELECT\\n        year,\\n        avg_horsepower,\\n        LAG(avg_horsepower) OVER (ORDER BY year) AS prev_year_hp,\\n        LEAD(avg_horsepower) OVER (ORDER BY year) AS next_year_hp\\n    FROM yearly_avg_hp\\n)\\n\\nSELECT\\n    year,\\n    avg_horsepower,\\n    ROUND((\\n        (\\n            avg_horsepower - prev_year_hp\\n        ) / NULLIF(prev_year_hp, 0) * 100\\n    ), 2) AS percentage_diff_prev_year,\\n    ROUND((\\n        (\\n            next_year_hp - avg_horsepower\\n        ) / NULLIF(avg_horsepower, 0) * 100\\n    ), 2) AS percentage_diff_next_year\\nFROM yearly_comparisons\\nORDER BY\\n    year\"}\n",
      ".."
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>YEAR</th>\n",
       "      <th>AVG_HORSEPOWER</th>\n",
       "      <th>PERCENTAGE_DIFF_PREV_YEAR</th>\n",
       "      <th>PERCENTAGE_DIFF_NEXT_YEAR</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1970</td>\n",
       "      <td>148.857143</td>\n",
       "      <td>NaN</td>\n",
       "      <td>-29.51</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1971</td>\n",
       "      <td>104.928571</td>\n",
       "      <td>-29.51</td>\n",
       "      <td>14.53</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1972</td>\n",
       "      <td>120.178571</td>\n",
       "      <td>14.53</td>\n",
       "      <td>8.57</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1973</td>\n",
       "      <td>130.475000</td>\n",
       "      <td>8.57</td>\n",
       "      <td>-27.78</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1974</td>\n",
       "      <td>94.230769</td>\n",
       "      <td>-27.78</td>\n",
       "      <td>7.25</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1975</td>\n",
       "      <td>101.066667</td>\n",
       "      <td>7.25</td>\n",
       "      <td>0.05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1976</td>\n",
       "      <td>101.117647</td>\n",
       "      <td>0.05</td>\n",
       "      <td>3.91</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1977</td>\n",
       "      <td>105.071429</td>\n",
       "      <td>3.91</td>\n",
       "      <td>-5.12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1978</td>\n",
       "      <td>99.694444</td>\n",
       "      <td>-5.12</td>\n",
       "      <td>1.52</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1979</td>\n",
       "      <td>101.206897</td>\n",
       "      <td>1.52</td>\n",
       "      <td>-23.44</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>1980</td>\n",
       "      <td>77.481481</td>\n",
       "      <td>-23.44</td>\n",
       "      <td>5.88</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>1981</td>\n",
       "      <td>82.034483</td>\n",
       "      <td>5.88</td>\n",
       "      <td>-0.69</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>1982</td>\n",
       "      <td>81.466667</td>\n",
       "      <td>-0.69</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    YEAR  AVG_HORSEPOWER  PERCENTAGE_DIFF_PREV_YEAR  PERCENTAGE_DIFF_NEXT_YEAR\n",
       "0   1970      148.857143                        NaN                     -29.51\n",
       "1   1971      104.928571                     -29.51                      14.53\n",
       "2   1972      120.178571                      14.53                       8.57\n",
       "3   1973      130.475000                       8.57                     -27.78\n",
       "4   1974       94.230769                     -27.78                       7.25\n",
       "5   1975      101.066667                       7.25                       0.05\n",
       "6   1976      101.117647                       0.05                       3.91\n",
       "7   1977      105.071429                       3.91                      -5.12\n",
       "8   1978       99.694444                      -5.12                       1.52\n",
       "9   1979      101.206897                       1.52                     -23.44\n",
       "10  1980       77.481481                     -23.44                       5.88\n",
       "11  1981       82.034483                       5.88                      -0.69\n",
       "12  1982       81.466667                      -0.69                        NaN"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got output: [Document(id_='2a53eaf3-ffa3-4ce0-9d7d-60575ae77f63', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='721222e7ca4249d174a35e130eb622c0c9a7b46211a27f7998008ee9aeff31a6', text=\"{'YEAR': 1970, 'AVG_HORSEPOWER': 148.857143, 'PERCENTAGE_DIFF_PREV_YEAR': None, 'PERCENTAGE_DIFF_NEXT_YEAR': -29.51}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='750de6d2-6872-4d48-bbce-65c6ae020c14', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='f8035cca9f3954c9b453303c9677871c23741ecff32d5070f30e579fa8d38e10', text=\"{'YEAR': 1971, 'AVG_HORSEPOWER': 104.928571, 'PERCENTAGE_DIFF_PREV_YEAR': -29.51, 'PERCENTAGE_DIFF_NEXT_YEAR': 14.53}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='271cb85e-fbaf-4e79-87b2-01203c3ae78d', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='b50d15eaa536ff504088f9469773ad8f4796a04ee0f7682b20047768031369e4', text=\"{'YEAR': 1972, 'AVG_HORSEPOWER': 120.178571, 'PERCENTAGE_DIFF_PREV_YEAR': 14.53, 'PERCENTAGE_DIFF_NEXT_YEAR': 8.57}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='b9307acb-96ec-465f-8acd-79a62b83a50e', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='6a2badb344e6028a7e5423b242d87225f1dc7b789e40740612d0706db9a94cbf', text=\"{'YEAR': 1973, 'AVG_HORSEPOWER': 130.475, 'PERCENTAGE_DIFF_PREV_YEAR': 8.57, 'PERCENTAGE_DIFF_NEXT_YEAR': -27.78}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='b573917c-7d05-47c3-bac7-5d0465d819c6', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='91c410d7850e434fc2279b12d9d5db7d39f8ed1f10bb5824b53ecc6b6a4f203e', text=\"{'YEAR': 1974, 'AVG_HORSEPOWER': 94.230769, 'PERCENTAGE_DIFF_PREV_YEAR': -27.78, 'PERCENTAGE_DIFF_NEXT_YEAR': 7.25}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='ef3afbd8-4169-4194-a526-c38fc965e87c', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='65db95a2beb9a1a769dd3b42b8a52a75c8e901bc3f10bd3525f97300f3154d29', text=\"{'YEAR': 1975, 'AVG_HORSEPOWER': 101.066667, 'PERCENTAGE_DIFF_PREV_YEAR': 7.25, 'PERCENTAGE_DIFF_NEXT_YEAR': 0.05}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='31f3f03c-4a8a-42e7-9514-8f552bb00838', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='e85f69f2bcd31b6a9f6c0c14d8393c500ac3d0cc790fd16a920da5b881646e7f', text=\"{'YEAR': 1976, 'AVG_HORSEPOWER': 101.117647, 'PERCENTAGE_DIFF_PREV_YEAR': 0.05, 'PERCENTAGE_DIFF_NEXT_YEAR': 3.91}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='b0cd75b7-2ce9-4a32-92b2-09322da5194b', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='2a30268f65fa21e25875b9e1b1998375287cccce38b2628110c5924dd674156e', text=\"{'YEAR': 1977, 'AVG_HORSEPOWER': 105.071429, 'PERCENTAGE_DIFF_PREV_YEAR': 3.91, 'PERCENTAGE_DIFF_NEXT_YEAR': -5.12}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='14c84aff-21a4-4abb-ba98-ad1643eb45bf', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='a6f93a216fdb4afd7a3c2215f1838eb551e331079ae604429f8dee8848f318ca', text=\"{'YEAR': 1978, 'AVG_HORSEPOWER': 99.694444, 'PERCENTAGE_DIFF_PREV_YEAR': -5.12, 'PERCENTAGE_DIFF_NEXT_YEAR': 1.52}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='11b59de3-1d01-4020-9cb6-48ecf1278ab7', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='c0b41a6866f0d5e14b81c67a5fb772ba6877643079f43d631734ce7b971bf6bf', text=\"{'YEAR': 1979, 'AVG_HORSEPOWER': 101.206897, 'PERCENTAGE_DIFF_PREV_YEAR': 1.52, 'PERCENTAGE_DIFF_NEXT_YEAR': -23.44}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='6c652dfe-2e12-432f-9dce-a3f3c098ea6e', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='7ff1cd7f6b343f65dc19194f125374edd874ebab6bbe0eb813ddd85b4cb8499a', text=\"{'YEAR': 1980, 'AVG_HORSEPOWER': 77.481481, 'PERCENTAGE_DIFF_PREV_YEAR': -23.44, 'PERCENTAGE_DIFF_NEXT_YEAR': 5.88}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='ddb554f2-e704-45d6-9864-80d3d7f1a319', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='d3977f6455afd966740cde1dc065ea132387ad3ce761b1b52dcddc881dc0f105', text=\"{'YEAR': 1981, 'AVG_HORSEPOWER': 82.034483, 'PERCENTAGE_DIFF_PREV_YEAR': 5.88, 'PERCENTAGE_DIFF_NEXT_YEAR': -0.69}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), Document(id_='298a6b49-ce19-463b-ab18-6d591d2aaca8', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='e0f5a233bdf0e2b4102d7aa1bcb25856eb32777c9b7ab5b9645e8b205972f545', text=\"{'YEAR': 1982, 'AVG_HORSEPOWER': 81.466667, 'PERCENTAGE_DIFF_PREV_YEAR': -0.69, 'PERCENTAGE_DIFF_NEXT_YEAR': None}\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')]\n",
      "========================\n",
      "\n",
      "STARTING TURN 2\n",
      "---------------\n",
      "\n",
      "The query has been executed, and here are the results for the average horsepower (`avg_horsepower`) by year, along with the percentage difference from the previous year (`percentage_diff_prev_year`) and the next year (`percentage_diff_next_year`):\n",
      "\n",
      "- **1970**: avg_horsepower: 148.86, percentage_diff_prev_year: None, percentage_diff_next_year: -29.51%\n",
      "- **1971**: avg_horsepower: 104.93, percentage_diff_prev_year: -29.51%, percentage_diff_next_year: 14.53%\n",
      "- **1972**: avg_horsepower: 120.18, percentage_diff_prev_year: 14.53%, percentage_diff_next_year: 8.57%\n",
      "- **1973**: avg_horsepower: 130.48, percentage_diff_prev_year: 8.57%, percentage_diff_next_year: -27.78%\n",
      "- **1974**: avg_horsepower: 94.23, percentage_diff_prev_year: -27.78%, percentage_diff_next_year: 7.25%\n",
      "- **1975**: avg_horsepower: 101.07, percentage_diff_prev_year: 7.25%, percentage_diff_next_year: 0.05%\n",
      "- **1976**: avg_horsepower: 101.12, percentage_diff_prev_year: 0.05%, percentage_diff_next_year: 3.91%\n",
      "- **1977**: avg_horsepower: 105.07, percentage_diff_prev_year: 3.91%, percentage_diff_next_year: -5.12%\n",
      "- **1978**: avg_horsepower: 99.69, percentage_diff_prev_year: -5.12%, percentage_diff_next_year: 1.52%\n",
      "- **1979**: avg_horsepower: 101.21, percentage_diff_prev_year: 1.52%, percentage_diff_next_year: -23.44%\n",
      "- **1980**: avg_horsepower: 77.48, percentage_diff_prev_year: -23.44%, percentage_diff_next_year: 5.88%\n",
      "- **1981**: avg_horsepower: 82.03, percentage_diff_prev_year: 5.88%, percentage_diff_next_year: -0.69%\n",
      "- **1982**: avg_horsepower: 81.47, percentage_diff_prev_year: -0.69%, percentage_diff_next_year: None\n",
      "\n",
      "Please note that the percentage differences for the first and last years are not available because there is no previous year data for 1970 and no next year data for 1982.\n"
     ]
    }
   ],
   "source": [
    "print(agent.chat(\"Run it\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b231bfe3-6c69-4458-a185-3e13db358e7d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myenv",
   "language": "python",
   "name": "myenv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
